from abpy import *
from time import sleep
from MyTT import *
import numpy as np
import os
from datetime import datetime
import pymssql
from sqlalchemy import create_engine, distinct, or_, and_
from urllib.parse import quote_plus as urlquote


def initialize(context):
    print("初始化")
    log.info(context.current_dt)

    # 外网地址
    # g.host = '218.17.227.194'
    # g.port = 10001

    # 内网地址
    g.host = '10.81.21.140'
    g.port = '1433'
    g.user = 'sa'
    g.password = 'Jzt@20230510&&'
    g.schema = 'dbo'
    g.db_name = 'AlphabeeCom'
    g.table_name = 'KUANDU_30MIN'
    g.endT = datetime.now().strftime('%Y%m%d') + '160000'
    g.index = dict()

    sql = f"select * from INDEX_CONS"
    conn = pymssql.connect(host=g.host, port=g.port, user=g.user, password=g.password, database=g.db_name)
    cursor = conn.cursor()
    cursor.execute(sql)
    r1 = cursor.fetchall()

    for tok in r1:
        index_id = tok[0]
        index_id = index_id + '.00' if index_id[:2] == '39' else index_id + '.10'
        code = tok[1]
        code = code + '.00' if code[:2] in ['00', '30'] else code + '.10'
        if index_id in g.index:
            g.index[index_id].append(code)
        else:
            g.index[index_id] = [code]

    # call_anytime()


def handle_data(context, data):
    print("enter handle_data")

    call_anytime()


def call_anytime():
    time_list = get_history_data('399300.00', 80, '30m', ['close'], end_dt=g.endT)['datetime']
    time_list.reverse()

    df = calDayValue_kuandu(time_list)
    print(df)
    # if len(df) > 0:
    #     insert_db(df)
    #     print(f'all index insert ok')
    # else:
    #     print(f'no index insert')


def calDayValue_kuandu(time_list):
    w7, w13, w30 = 7, 13, 30
    SN, MN, LN, VN = 7, 13, 21, 35
    period = 8
    ret = list()
    # ret_n = get_update_n()
    ret_n = {'000016.10': 88, '000688.10': 88, '000852.10': 88, '000905.10': 88, '399004.00': 88, '399006.00': 88, '399300.00': 88, '399330.00': 88}
    print(f'ret_n={ret_n}')

    for index in g.index:
        HS1, LS1, HS2, LS2, HS3, LS3 = 0, 0, 0, 0, 0, 0
        SS_AD, SS_DI, MS_AD, MS_DI, LS_AD, LS_DI, VS_AD, VS_DI = 0, 0, 0, 0, 0, 0, 0, 0
        print(index)
        n = ret_n[index]

        if n > 0:
            stockArr = g.index[index]
            for code in stockArr:
                # print(f'code={code}')
                data = get_history_data(code, VN + 1 + 120, '30m', ['high', 'low', 'close'], end_dt=g.endT)
                # data = get_data_last_minute(data)
                # print(f'code={code}->data={data}')
                high_list = data['high']
                low_list = data['low']
                close_list = data['close']
                high_list.reverse()
                low_list.reverse()
                close_list.reverse()
                high_list = np.array(high_list)
                low_list = np.array(low_list)
                close_list = np.array(close_list)

                vH7, vL7 = REF(HHV(high_list, w7), 1), REF(LLV(low_list, w7), 1)
                vH13, vL13 = REF(HHV(high_list, w13), 1), REF(LLV(low_list, w13), 1)
                vH30, vL30 = REF(HHV(high_list, w30), 1), REF(LLV(low_list, w30), 1)
                # print(code.center(20, '*'))
                # print(f'high_list={high_list}->low_list={low_list}')
                # print(f'vH7={vH7}->vL7={vL7}')
                # print(f'vH13={vH13}->vL13={vL13}')
                # print(f'vH30={vH30}->vL30={vL30}')

                #      print(len(high_list), len(low_list), len(close_list))
                HS1 += (high_list > vH7)
                LS1 += (low_list < vL7)
                HS2 += (high_list > vH13)
                LS2 += (low_list < vL13)
                HS3 += (high_list > vH30)
                LS3 += (low_list < vL30)

                cc1 = SMA(close_list, SN, 1)
                cc2 = SMA(close_list, MN, 1)
                cc3 = SMA(close_list, LN, 1)
                cc4 = SMA(close_list, VN, 1)

                SS_AD += (divide(close_list, cc1) - 1 > 0)
                SS_DI += (divide(close_list, cc1) - 1 <= 0)
                MS_AD += (divide(close_list, cc2) - 1 > 0)
                MS_DI += (divide(close_list, cc2) - 1 <= 0)
                LS_AD += (divide(close_list, cc3) - 1 > 0)
                LS_DI += (divide(close_list, cc3) - 1 <= 0)
                VS_AD += (divide(close_list, cc4) - 1 > 0)
                VS_DI += (divide(close_list, cc4) - 1 <= 0)

            HS1 = HS1[w7:]
            LS1 = LS1[w7:]
            HS2 = HS2[w13:]
            LS2 = LS2[w13:]
            HS3 = HS3[w30:]
            LS3 = LS3[w30:]
            # print(f'HS1={HS1[-16:]}->LS1={LS1[-16:]}')
            # print(f'HS2={HS2[-16:]}->LS2={LS2[-16:]}')
            # print(f'HS3={HS3[-16:]}->LS3={LS3[-16:]}')
            A1 = divide(HS1, HS1 + LS1) * 100
            A2 = divide(HS2, HS2 + LS2) * 100
            A3 = divide(HS3, HS3 + LS3) * 100
            SS_AD = SS_AD[SN:]
            SS_DI = SS_DI[SN:]
            MS_AD = MS_AD[MN:]
            MS_DI = MS_DI[MN:]
            LS_AD = LS_AD[LN:]
            LS_DI = LS_DI[LN:]
            VS_AD = VS_AD[VN:]
            VS_DI = VS_DI[VN:]
            LXS = (SS_AD - SS_DI)
            LXM = (MS_AD - MS_DI)
            LXL = (LS_AD - LS_DI)
            LXV = (VS_AD - VS_DI)

            KD1 = EMA(A1, period)[-n:]
            KD2 = EMA(A2, period)[-n:]
            KD3 = EMA(A3, period)[-n:]
            KS1 = EMA(LXS, SN)[-n:]
            KS2 = EMA(LXM, MN)[-n:]
            KS3 = EMA(LXL, LN)[-n:]
            KS4 = EMA(LXV, VN)[-n:]
            time_list = time_list[-n:]

            for i in range(n):
                dt = time_list[i]
                dt = int(dt.strftime('%Y%m%d%H%M'))
                item = [index[:6], dt, KD1[i], KD2[i], KD3[i], KS1[i], KS2[i], KS3[i], KS4[i]]
                ret.append(item)

    df = pd.DataFrame(ret, columns=['STK_CODE', 'DATE_T', 'KD1', 'KD2', 'KD3',
                                    'KS1', 'KS2', 'KS3', 'KS4'])
    # print(df)

    return df


def get_update_n():
    conn = pymssql.connect(host=g.host, port=g.port, user=g.user, password=g.password, database=g.db_name)
    cursor = conn.cursor()
    sql = f"SELECT STK_CODE, MAX(DATE_T) FROM {g.table_name} GROUP BY STK_CODE;"
    cursor.execute(sql)
    r = cursor.fetchall()
    time_list = get_history_data('399300.00', 160, '30m', ['close'], end_dt=g.endT)['datetime']
    ret = dict()
    for tok in r:
        index_id = tok[0]
        index_id = index_id + '.00' if index_id[:2] == '39' else index_id + '.10'
        ts = str(tok[1])
        dt = datetime.strptime(ts, '%Y%m%d%H%M')
        if dt in time_list:
            n = time_list.index(dt)
            ret[index_id] = n
    return ret


def get_data_last_minute(all_data):
    t = get_time()
    cols = all_data.keys()
    date_list = all_data['datetime']
    # print(f'all_data1={all_data}')
    idx = 0
    flag = False

    if t in date_list:
        flag = True
        idx = date_list.index(t)

    res = dict()

    for col in cols:
        res[col] = all_data[col][idx:]

    # print(f"get_data_last_minute->flag={flag}")
    # print(f'all_data2={all_data}')

    return res


def get_time():
    curtime = datetime.now()
    year = curtime.year
    month = curtime.month
    day = curtime.day
    hour = curtime.hour
    minute = curtime.minute

    #    t = datetime.strptime(f'{year}-{month}-{day} {hour}:{minute}:0', '%Y-%m-%d %H:%M:%S')
    t = datetime.strptime(f'{year}-{month}-{day - 1} 15:0:0', '%Y-%m-%d %H:%M:%S')
    return t


def insert_db(df):
    yconnect = create_engine(
        f'mssql+pymssql://{g.user}:{urlquote(g.password)}@{g.host}:{g.port}/{g.db_name}?charset=utf8')
    pd.io.sql.to_sql(df, g.table_name, yconnect, schema=g.schema, if_exists='append', index=False)


def divide(lst1, lst2):
    res = []
    for i in range(len(lst1)):
        a = lst1[i]
        b = lst2[i]
        r = 0 if b == 0 else a / b
        res.append(r)

    return np.array(res)
